from load_config import ConfigLoader
from Dataloaders.loadASCAD import Datasetloader
from Trainer.supervised_trainer import SupervisedTrainer
from Trainer.unsupervised_trainer import UnsupervisedTrainer
from Trainer.ddla_trainer import DDLATrainer
from Valer.valer import Valer
import os

config_loader = ConfigLoader()
configs = config_loader.get_configs()
for i, config in enumerate(configs):
    print(f"{i+1}/{len(configs)}. 加载配置......")
    print(f"配置详情：{config}")
    # 新建项目所有文件夹
    if not os.path.exists("./project_datasave"):
        os.mkdir("./project_datasave")
    if not os.path.exists(f"./project_datasave/{config['project_name']}"):
        os.mkdir(f"./project_datasave/{config['project_name']}")
        os.mkdir(f"./project_datasave/{config['project_name']}/model")
        os.mkdir(f"./project_datasave/{config['project_name']}/figure")
        os.mkdir(f"./project_datasave/{config['project_name']}/log")
    # 加载数据集
    if (config["trainer"]["type"] != "ddla"):
        data_loader = Datasetloader(config)
        train_loader = data_loader.get_train_loader()
        val_loader = data_loader.get_val_loader()
    # 训练
    if config["trainer"]["type"] == "supervised":
        trainer = SupervisedTrainer(config, train_loader)
    elif config["trainer"]["type"] == "unsupervised":
        trainer = UnsupervisedTrainer(config, train_loader)
    elif config["trainer"]["type"] == "ddla":
        trainer = DDLATrainer(config)

    trainer.train()
    print("训练完成，训练时间：", trainer.train_time)
    if (config["trainer"]["type"] == "ddla"):
        print("分组时间：", trainer.group_time)
    # 验证（攻击）
    if (config["trainer"]["type"] != "ddla"):
        valer = Valer(config, val_loader)
        valer.val()
        print("验证完成，验证准确率时间：", valer.val_time)
        print("攻击时间：", valer.attack_time)
